# -*- coding: utf-8 -*-
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import gaussian_kde
from matplotlib.patches import Ellipse
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from pyproj import Transformer
import geopandas as gpd
from shapely.geometry import Point
import matplotlib.gridspec as gridspec

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


# ==================== 1. 数据加载与预处理 ====================
def load_data():
    """加载并清洗西安POI数据"""
    try:
        df = pd.read_csv('warehouse_xian_poi.csv')
        print(f"原始数据量: {len(df)}条")

        # 数据清洗
        df = df.dropna(subset=['lng', 'lat', 'query_category'])
        print(f"去除空值后: {len(df)}条")

        # 筛选旅游相关类别
        tourism_types = ['景点', '文物古迹', '博物馆', '公园', '游乐园',
                         '寺庙', '植物园', '动物园', '水族馆', '教堂']
        df_tourism = df[df['query_category'].isin(tourism_types)].copy()
        print(f"筛选旅游类别后: {len(df_tourism)}条")

        # 转换为GeoDataFrame
        geometry = [Point(xy) for xy in zip(df_tourism['lng'], df_tourism['lat'])]
        gdf = gpd.GeoDataFrame(df_tourism, geometry=geometry, crs="EPSG:4326")

        return gdf
    except Exception as e:
        print(f"数据加载出错: {str(e)}")
        return gpd.GeoDataFrame()


# ==================== 2. 标准差椭圆计算 ====================
def calculate_ellipse_params(gdf):
    """计算各类别标准差椭圆参数"""
    if gdf.empty:
        return pd.DataFrame()

    # 坐标转换（WGS84转UTM Zone 49N）
    transformer = Transformer.from_crs(4326, 32649)

    def _ellipse(group):
        if len(group) < 2:
            return pd.Series(dtype=float)

        # 转换为平面坐标
        coords = np.array([transformer.transform(point.y, point.x)
                           for point in group.geometry])
        x, y = coords[:, 0], coords[:, 1]

        # 计算椭圆参数
        cov = np.cov(x - x.mean(), y - y.mean())
        eigvals, eigvecs = np.linalg.eig(cov)
        major = 2 * np.sqrt(eigvals[0])
        minor = 2 * np.sqrt(eigvals[1])
        angle = np.degrees(np.arctan2(*eigvecs[:, 0][::-1])) % 180

        # 计算中心点
        center = group.geometry.union_all().centroid

        return pd.Series({
            'center_lng': center.x,
            'center_lat': center.y,
            'major_axis': major,
            'minor_axis': minor,
            'angle': angle,
            'area': np.pi * major * minor / 1e6,  # 转换为平方公里
            'count': len(group)
        })

    return gdf.groupby('query_category', group_keys=False).apply(_ellipse)


# ==================== 3. 可视化函数 ====================
def plot_combined_map(gdf, ellipse_params):
    """在西安地图上同时显示标准差椭圆和核密度"""
    if gdf.empty:
        return

    # 创建地图画布
    plt.figure(figsize=(16, 12))
    ax = plt.axes(projection=ccrs.PlateCarree())

    # 设置西安地图范围（包含周边区域）
    ax.set_extent([108.7, 109.3, 33.9, 34.5])  # 扩大显示范围

    # 添加专业地图元素
    ax.add_feature(cfeature.LAND.with_scale('50m'), facecolor='#f5f5f5')
    ax.add_feature(cfeature.RIVERS.with_scale('50m'), edgecolor='#66b3ff', linewidth=1)
    ax.add_feature(cfeature.LAKES.with_scale('50m'), edgecolor='#66b3ff', facecolor='#a6cee3')
    ax.gridlines(draw_labels=True, linestyle='--', alpha=0.5)

    # ===== 1. 绘制核密度 =====
    # 计算整体核密度
    kde = gaussian_kde(np.vstack([gdf.geometry.x, gdf.geometry.y]))
    xi, yi = np.mgrid[108.7:109.3:200j, 33.9:34.5:200j]
    zi = kde(np.vstack([xi.flatten(), yi.flatten()]))

    # 绘制核密度热力图（使用半透明紫色调）
    levels = np.linspace(zi.min(), zi.max(), 10)
    cs = ax.contourf(xi, yi, zi.reshape(xi.shape), levels=levels,
                     cmap='Purples', alpha=0.5, transform=ccrs.PlateCarree())

    # ===== 2. 绘制标准差椭圆 =====
    colors = plt.cm.tab20(np.linspace(0, 1, len(ellipse_params)))
    for idx, (category, row) in enumerate(ellipse_params.iterrows()):
        # 绘制椭圆
        ellipse = Ellipse(
            (row['center_lng'], row['center_lat']),
            width=row['minor_axis'] / 50000,
            height=row['major_axis'] / 50000,
            angle=row['angle'],
            edgecolor=colors[idx],
            facecolor='none',
            linewidth=2,
            transform=ccrs.PlateCarree(),
            label=f"{category} (n={int(row['count'])})"
        )
        ax.add_patch(ellipse)

        # 标记中心点
        ax.plot(row['center_lng'], row['center_lat'], 'o',
                color=colors[idx], markersize=8, transform=ccrs.PlateCarree())

    # 添加图例和标题
    plt.legend(title='景点类别(数量)', loc='upper right', fontsize=9)
    plt.title('西安市旅游景点空间分布\n(核密度与标准差椭圆)', pad=20, fontsize=16)

    # 添加比例尺和指北针
    def add_scale_bar():
        ax.plot([108.75, 108.85], [33.92, 33.92], color='black', linewidth=2,
                transform=ccrs.PlateCarree())
        ax.text(108.8, 33.91, '10 km', ha='center', transform=ccrs.PlateCarree())

    def add_north_arrow():
        ax.annotate('↑ N', xy=(0.95, 0.95), xycoords='axes fraction',
                    ha='center', va='center', fontsize=12,
                    bbox=dict(boxstyle='round', facecolor='white'))

    add_scale_bar()
    add_north_arrow()

    # 添加颜色条
    cbar = plt.colorbar(cs, orientation='vertical', shrink=0.6, pad=0.02)
    cbar.set_label('景点分布密度', rotation=270, labelpad=15)

    plt.savefig('xian_tourism_combined_map.png', dpi=300, bbox_inches='tight')
    plt.show()


# ==================== 主执行流程 ====================
if __name__ == "__main__":
    # 1. 数据加载
    gdf = load_data()

    if not gdf.empty:
        # 2. 计算标准差椭圆参数
        ellipse_params = calculate_ellipse_params(gdf)
        print("\n标准差椭圆参数表:")
        print(ellipse_params.round(1))

        # 3. 绘制综合地图
        plot_combined_map(gdf, ellipse_params)

        # 4. 保存结果
        ellipse_params.to_excel('xian_ellipse_params.xlsx')
    else:
        print("执行失败：无有效数据")